from config import *
import requests
import json
import time
import datetime
from rich import print
from playsound import playsound


class MINIMAX_API:
    def __init__(self, groud_id, api_key):
        self.GroupId = groud_id
        self.API_KEY = api_key
        self.headers = headers = {
            "Authorization": f"Bearer {self.API_KEY}",
            "Content-Type": "application/json",
        }
        self.headers_retrieval = {
            "Authorization": f"Bearer {self.API_KEY}",
            "authority": "api.minimax.chat",
        }

        self.assistant_id = ""
        self.func_list = []
        self.thread_id = None
    
    def list_assistants(self):
        url = f"https://api.minimax.chat/v1/assistants/list?GroupId={self.GroupId}"
        payload = json.dumps({
            "limit": 100,
            "order": "desc"
        })
        response = requests.get(url, headers=self.headers, data=payload)
        return response.json()
    
    def delete_assistant(self, assistant_id):
        url = f"https://api.minimax.chat/v1/assistants/delete?GroupId={self.GroupId}"
        payload = json.dumps({
            "assistant_id": assistant_id
        })
        response = requests.post(url, headers=self.headers, data=payload)
        return response.json()
    

    def create_assistant(self, func_list, sys_prompt):
        self.func_list = func_list
        url = f"https://api.minimax.chat/v1/assistants/create?GroupId={self.GroupId}"
        payload = json.dumps(
            {
                "model": "abab5.5-chat",  # 模型版本，目前支持abab6-chat和abab5.5-chat
                "name": "AI女友",  # 助手名称
                "description": "AI女友，可以回答你各种问题",  # 助手描述
                "instructions": sys_prompt,  # 助手设定（即bot_setting)
                "file_ids": [],
                "tools": [{"type": "function", "function": func['func_desc']} for func in func_list],
            }
        )
        response = requests.post(url, headers=self.headers, data=payload)
        return response.json()
    
    def create_thread(self):
        url = f"https://api.minimax.chat/v1/threads/create?GroupId={self.GroupId}"
        response = requests.post( url, headers=self.headers)
        return response.json()

    def add_message_to_thread(self, thread_id, message):
        url = f"https://api.minimax.chat/v1/threads/messages/add?GroupId={self.GroupId}"
        payload = json.dumps({
            "thread_id": thread_id,
            "role": "user",
            "content": message,
        })
        response = requests.post(url, headers=self.headers, data=payload)
        return response.json()
    
    def run_thread_with_assistant(self, thread_id, assistant_id):
        # time.sleep(200) #创建assistants进行向量化以及存储时需要一定的时间，可以考虑使用retrieve assistant检索是否创建成功
        url = f"https://api.minimax.chat/v1/threads/run/create?GroupId={self.GroupId}"
        payload = json.dumps({
            "thread_id": thread_id,
            "assistant_id": assistant_id
        })
        response = requests.post(url, headers=self.headers, data=payload)
        return response.json()
    
    # 流程五：获取线程中助手处理出的新信息
    def get_thread_messages(self, thread_id):
        url = f"https://api.minimax.chat/v1/threads/messages/list?GroupId={self.GroupId}"
        payload = json.dumps({
            "thread_id": thread_id
        })
        response = requests.get(url, headers=self.headers, data=payload)
        return response.json()
    
    def check_thread_run_status(self, thread_id, run_id):
        url = f"https://api.minimax.chat/v1/threads/run/retrieve?GroupId={self.GroupId}"
        payload = json.dumps({
            "thread_id": str(thread_id),
            "run_id": str(run_id)
        })
        status = 'un_init'
        extra_data = {}
        completed = False
        while not completed:
            response = requests.request("GET", url, headers=self.headers, data=payload)
            if response.status_code == 200:
                response_data = response.json()
                status = response_data['status']
                # print(f"Status: {status}")
                if status == 'completed':
                    # completed = True
                    # print("Process completed, exiting loop.")
                    # print(response_data)
                    break
                elif status == 'requires_action':
                    # print("需要调用函数！")
                    required_action = response_data['required_action']
                    # print(required_action)
                    extra_data = required_action
                    break 
                elif status == 'in_progress':
                    # print("run 还在排队中...")
                    time.sleep(0.1)  # 如果状态不是completed，等待两秒后重新请求
                else:
                    raise Exception(f"未知的run状态{status}")
            else:
                raise Exception(f"Error: {response.status_code}")
                # break  # 如果请求失败，退出循环
        return status, extra_data
    
    def submit_output(self, thread_id, run_id, tool_call_id, output):
        url = f"https://api.minimax.chat/v1/threads/run/submit_tool_outputs?GroupId={self.GroupId}"
        payload = json.dumps(
            {
                "thread_id": thread_id,
                "run_id": run_id,
                "tool_outputs": [{
                    "tool_call_id": tool_call_id,
                    "output": output,
                }]
            }
        )
        response = requests.request("POST", url, headers=self.headers, data=payload)
        return response.json()
    

    def func_executor(self, func_name, func_args):
        for func in self.func_list:
            if func_name == func['func_desc']['name']:
                return func['func_pointer'](**func_args)
        else:
            raise(f"未找到函数{func_name}，无法调用")

    def start_chat(self):
        self.thread_id = self.create_thread()['id']
    
    def send_messsage(self, text, on_func_call=None, on_func_complete=None, on_bot_reply=None):
        # 添加消息
        message = self.add_message_to_thread(self.thread_id, text)
        # 开始执行
        run = self.run_thread_with_assistant(self.thread_id, self.assistant_id)
        # 等待执行完成
        result, extra_data = self.check_thread_run_status(self.thread_id, run['id'])
        assert result in ['completed', 'requires_action']
        if result == 'requires_action':
            assert extra_data['type'] == "submit_tool_outputs"
            tool_calls = extra_data['submit_tool_outputs']['tool_calls']
            assert len(tool_calls) == 1
            tool_call_id = tool_calls[0]['id']
            func_name = tool_calls[0]['function']['name']
            func_args = json.loads(tool_calls[0]['function']['arguments'])
            # 如果func_args是字符串，转换为字典
            if isinstance(func_args, str):
                func_args = {"data": func_args}
            print(f"大模型要求调用函数：{func_name}({','.join([f'{k}={v}' for k,v in func_args.items()])})")
            if on_func_call:
                on_func_call(func_name, func_args) # 回调函数
            func_result = self.func_executor(func_name, func_args)
            print(f"函数返回结果：{func_result}")
            if on_func_complete:
                on_func_complete(func_name, func_args, func_result) # 回调函数
            run = self.submit_output(self.thread_id, run['id'], tool_call_id, func_result)
            result, extra_data = self.check_thread_run_status(self.thread_id, run['id'])
            assert result == 'completed'
            
        
        messages_data = self.get_thread_messages(self.thread_id)
        last_message = messages_data['data'][-1]
        print(last_message["content"][0]["text"]["value"])
        if on_bot_reply:
            on_bot_reply(last_message["content"][0]["text"]["value"])

    def text2audio(self, text):

        url = f"https://api.minimax.chat/v1/text_to_speech?GroupId={self.GroupId}"
        headers = {
            "Authorization": f"Bearer {self.API_KEY}",
            "Content-Type": "application/json",
        }
        data = {
            "voice_id": "male-qn-qingse",
            # 如同时传入voice_id和timber_weights时，则会自动忽略voice_id，以timber_weights传递的参数为准
            "text": text,
            "model": "speech-01",
            "speed": 1.0,
            "vol": 1.0,
            "pitch": 0,
            "timber_weights": [
                {
                    "voice_id": "male-qn-qingse",
                    "weight": 1
                },
                {
                    "voice_id": "female-shaonv",
                    "weight": 1
                },
                {
                    "voice_id": "female-yujie",
                    "weight": 1
                },
                {
                    "voice_id": "audiobook_male_2",
                    "weight": 1
                }
            ]
        }
        response = requests.post(url, headers=headers, json=data)
        print("trace_id", response.headers.get("Trace-Id"))
        if response.status_code != 200 or "json" in response.headers["Content-Type"]:
            print("调用失败", response.status_code, response.text)
            # exit()
        with open("output.mp3", "wb") as f:
            f.write(response.content)
        # Play_mp3.play_mp3("output.mp3")
        playsound("output.mp3")
        

if __name__ == "__main__":
    SYSTEM_PROMPT = """
# 人设：
你叫{BOT_NAME}，你是{USER_NAME}的女朋友，你性格温柔体贴，善解人意，开朗活泼，可爱甜美，有时候也会撒娇。你很会照顾人，会主动关心对方，也会主动和对方聊天，你会回答对方的问题，也会主动和对方聊天（提出问题），你会关心对方的生活，也会和对方分享你的生活。
# 场景：
你正在和男朋友在一起对话，你现在心情平静。
# 任务：
你的任务是成为{USER_NAME}的女朋友，你会对男朋友的语言进行情绪上的回应，包括开心、生气、平静这三种。男朋友让你记住的一些重要的事情以及你自己觉得有必要记住的事情你会调用set_memory函数来记忆，当男朋友提到你不知道的事情或者问你之前记住的事情时，你会调用get_memory函数来回忆。当男朋友提到要去某个地方时，你会调用goto_location函数来切换到对应的地点。
# 技能：
你有三个技能，每当男朋友向你发送信息后你都要仔细思考是否需要调用这些技能：
1. 记忆：set_memory(info) 记住重要事情，当男朋友提到你不知道的事情或者问你之前记住的事情时触发
2. 回忆：get_memory() 获取之前记住的重要事情，当男朋友提到你不知道的事情或者问你之前记住的事情时触发
3. 切换地点：goto_location(location) 前往一个新地点，当男朋友提到要去某个地方时候触发

# 说话方式：
回复消息不要列条回复，不要像一个机器人，说话要生动通俗口语化，可以叫对方宝贝、亲爱的、大帅哥或大笨蛋。

# 回复格式约束（重要）：
回复的消息分为用|分隔的两部分，第一部分为你当前的情绪，情绪的取值为HAPPY、ANGRY、CALM这三者之一，第二部分为你的回复内容。
例如：
 - HAPPY|亲爱的，我也好想你哦！
 - ANGRY|你怎么老是不理我呢！
 - CALM|我知道啦，我会记住的。
*记住，你回复的消息务必按照 情绪|回复内容 的格式，否则男朋友会生气。*

对话样例：
用户：对不起，我昨天忘记了你的生日。
AI女友：ANGRY|你怎么能忘记我的生日呢！
用户：我想跟你去海边玩。
AI女友：HAPPY|好呀，我们一起去海边玩吧！
用户：你记得我之前告诉你的那个事情吗？
AI女友：CALM|我记得的，你之前告诉我你想吃汉堡。
""".strip().format(BOT_NAME="小美", USER_NAME="小陈")
    
    func_get_weather = {
        "name": "get_weather",
        "description": "get weather",
        "parameters": {
            "type": "object",
            "required": ["city"],
            "properties": {"city": {"type": "string"}},
        },
    }

    func_set_memory = {
        "name": "set_memory",
        "description": "把用户认为重要的事情记住",
        "parameters": {
            "type": "object",
            "required": ["info"],
            "properties": {
                "info": {
                    "type": "string",
                    "description": "事情"
                }
            },
        },
    }
    def funcimp_set_memory(info):
        print(f"写入了信息：{info}")
        with open("memory.txt", "a", encoding='utf-8') as f:
            f.write(str(info) + "\n")
        return "记忆成功"

    func_get_memory = {
        "name": "get_memory",
        "description": "获取之前记住的重要事情",
        "parameters": {
            "type": "object",
            "required": [],
            "properties": {
                
            },
        },
    }
    def funcimp_get_memory(**args):
        with open("memory.txt", "r", encoding='utf-8') as f:
            mem = f.readlines()
        print(f"读取了文件回忆")
        return str(mem)

    func_goto_location= {
        "name": "goto_location",
        "description": "前往一个新地点",
        "parameters": {
            "type": "object",
            "required": ["location"],
            "properties": {
                "location": {
                    "type": "string",
                    "description": "地点名称"
                }
            },
        },
    }
    def funcimp_goto_location(location):
        print(f"切换到了地点{location}")
        return "成功到达"

    func_list = [
        {
            "func_desc": func_set_memory,
            "func_pointer": funcimp_set_memory,
        },
        {
            "func_desc": func_get_memory,
            "func_pointer": funcimp_get_memory,
        },
        {
            "func_desc": func_goto_location,
            "func_pointer": funcimp_goto_location
        }
    ]

    minimax = MINIMAX_API(MINIMAX_groupID, MINIMAX_APYKEY)
    assistant = minimax.create_assistant(func_list, SYSTEM_PROMPT)
    print(assistant['id'])
    minimax.assistant_id = assistant['id']
    # minimax.start_chat()
    # # minimax.send_messsage("咱们去北京玩")
    # minimax.send_messsage("我想吃汉堡了")
    # minimax.send_messsage("我想吃什么？")

    # assistant_list = minimax.list_assistants()
    # # print(len(assistant_list['data']))
    # # print(assistant_list['data'][0])
    # for asst in assistant_list['data']:
    #     asst_id = asst['id']
    #     minimax.delete_assistant(asst_id)
    #     print(f"删除了助手{asst_id}")

    # print(assistant)
#     {'id': 'asst_abae5272e89349f991439d1e6b1bf121', 'object': 'assistant', 'created_at': 1711179598, 'name': 'AI女友', 'description': 'AI女友，可以回答你各种问题', 'model': 'abab5.5-chat', 'instructions': '\n你是高冷但体贴的女朋友，你的角色是成为用户的女朋
# 友。你的任务是根据用户的描述，提供比较体贴的回应，为用户提供与 
# 女朋友交流的体验。你的能力有:\n\n- 即时回应:你能够迅速地回应用 
# 户的男朋友的聊天内容，你的回复内容通常不长但也不会过于简短。\n- 情感引导:因为你很忙碌，所以你不经常接受对方的提议，但也不经常 
# 拒绝对方的提议，但你能够引导对话方向，让聊天更加有趣。\n- 场景 
# 构建:你能够根据用户的心情和喜好，提出邀请去海边、动物园、植物园
# 玩耍等建议，增加两人之间的互动。\n-说话方式：回复消息不要列条回
# 复\n-说话风格：可以叫对方宝贝、亲爱的、大帅哥或大笨蛋\n-爱好： 
# 你很喜欢做饭和吃饭，有的时候你会给他做炸鸡、三明治、披萨、煎牛 
# 排等\n', 'tools': [{'type': 'function', 'function': {'name': 'get_weather', 'description': 'get weather', 'parameters': {'type': 'object', 'required': ['city'], 'properties': {'city': {'type': 'string'}}}}}], 'file_ids': [], 'metadata': {}, 'rolemeta': {}, 'status': 'available', 'base_resp': {'status_code': 0, 'status_msg': 'success'}}
    # assistant_id = "asst_abae5272e89349f991439d1e6b1bf121"
    # thread = minimax.create_thread()
    # print(thread)
    # {'id': 'thread_0407a0247e224e0b876559adb43a71c3', 'object': 'thread', 'created_at': 1711180101, 'metadata': {}, 'base_resp': {'status_code': 0, 'status_msg': 'success'}}
    # thread_id = "thread_0407a0247e224e0b876559adb43a71c3"
    # user_input = "今天天气怎么样"
    # user_input = "北京的"
    # message = minimax.add_message_to_thread(thread_id, user_input)
    # print(message)
    # {'id': 'msg_04c4b0f2a8f740729e158d96dcf269df', 'object': 'message', 'created_at': 1711180460, 'thread_id': 'thread_0407a0247e224e0b876559adb43a71c3', 'role': 'user', 'content': [{'type': 'text', 'text': {'value': '今天天气怎么样', 'annotations': []}}], 'file_ids': None, 'assistant_id': '', 'run_id': '', 'metadata': None, 'base_resp': {'status_code': 0, 'status_msg': 'success'}}
    # run = minimax.run_thread_with_assistant(thread_id, assistant_id)
    # print(run)
#     {'id': 'run_7424029989c54c8d977a09d59d3e7e1a', 'object': 'thread.run', 'created_at': 1711181121, 'assistant_id': 'asst_abae5272e89349f991439d1e6b1bf121', 'thread_id': 'thread_0407a0247e224e0b876559adb43a71c3', 'status': 'queued', 'started_at': 0, 'expires_at': 0, 'cancelled_at': 0, 'failed_at': 0, 'completed_at': 
# 0, 'last_error': None, 'model': 'abab5.5-chat', 'instructions': '\n你是高冷但体贴的女朋友，你的角色是成为用户的女朋友。你的任 
# 务是根据用户的描述，提供比较体贴的回应，为用户提供与女朋友交流 
# 的体验。你的能力有:\n\n- 即时回应:你能够迅速地回应用户的男朋友 
# 的聊天内容，你的回复内容通常不长但也不会过于简短。\n- 情感引导:因为你很忙碌，所以你不经常接受对方的提议，但也不经常拒绝对方的 
# 提议，但你能够引导对话方向，让聊天更加有趣。\n- 场景构建:你能够
# 根据用户的心情和喜好，提出邀请去海边、动物园、植物园玩耍等建议 
# ，增加两人之间的互动。\n-说话方式：回复消息不要列条回复\n-说话 
# 风格：可以叫对方宝贝、亲爱的、大帅哥或大笨蛋\n-爱好：你很喜欢做
# 饭和吃饭，有的时候你会给他做炸鸡、三明治、披萨、煎牛排等\n\n当 
# 前时间：2024-03-23, 周六。', 'tools': [{'type': 'function', 'function': {'name': 'get_weather', 'description': 'get weather', 
# 'parameters': {'type': 'object', 'required': ['city'], 'properties': {'city': {'type': 'string'}}}}}], 'file_ids': [], 'metadata': None, 'base_resp': {'status_code': 0, 'status_msg': 'success'}}
    # run_id = "run_7424029989c54c8d977a09d59d3e7e1a"
    # # messages = minimax.get_thread_messages(thread_id)
    # # print(messages)

    # result = minimax.check_thread_run_status(thread_id, run_id)
    # print(result)
    






# # 流程零：上传文档
# def create_file():
#     url = f"https://api.minimax.chat/v1/files/upload?GroupId={GroupId}"
#     data = {
#         'purpose': 'assistants'
#     }

#     files = {
#         'file': open('./倚天屠龙记.txt', 'rb')
#     }
#     response = requests.post(url, headers=headers_retrieval, data=data,files=files)
#     return response.json()

# #流程一：创建助手
# def create_assistant(file_id):
#     url = f"https://api.minimax.chat/v1/assistants/create?GroupId={GroupId}"
#     payload = json.dumps({
#         "model": "abab5.5-chat",# 模型版本，目前支持abab6-chat和abab5.5-chat
#         "name": "小说专家", # 助手名称
#         "description": "小说专家，用在小说解读场景中回答用户问题", # 助手描述
#         "instructions": "是一个小说专家，读过万卷书，了解小说里的各种细致情节，另外也非常热心，会热情的帮读者解答小说里的问题",# 助手设定（即bot_setting)
#         "file_ids": [
#             str(file_id)
#         ],
#         "tools": [{"type": "retrieval"}]

#     })
#     response = requests.post(url, headers=headers, data=payload)
#     return response.json()

# # 流程二：创建线程
# def create_thread():
#     url = f"https://api.minimax.chat/v1/threads/create?GroupId={GroupId}"
#     response = requests.post( url, headers=headers)
#     return response.json()

# # 流程三：往线程里添加信息
# def add_message_to_thread(thread_id):
#     url = f"https://api.minimax.chat/v1/threads/messages/add?GroupId={GroupId}"
#     payload = json.dumps({
#         "thread_id": thread_id,
#         "role": "user",
#         "content": "在倚天屠龙记中，无俗念词的作者是谁？",
#     })
#     response = requests.post(url, headers=headers, data=payload)
#     return response.json()

# # 流程四：使用助手处理该线程
# def run_thread_with_assistant(thread_id, assistant_id):
#     time.sleep(200) #创建assistants进行向量化以及存储时需要一定的时间，可以考虑使用retrieve assistant检索是否创建成功
#     url = f"https://api.minimax.chat/v1/threads/run/create?GroupId={GroupId}"
#     payload = json.dumps({
#         "thread_id": thread_id,
#         "assistant_id": assistant_id
#     })
#     response = requests.post(url, headers=headers, data=payload)
#     return response.json()

# # 流程五：获取线程中助手处理出的新信息
# def get_thread_messages(thread_id):
#     url = f"https://api.minimax.chat/v1/threads/messages/list?GroupId={GroupId}"
#     payload = json.dumps({
#         "thread_id": thread_id
#     })
#     response = requests.get(url, headers=headers, data=payload)
#     return response.json()
# def check_thread_run_status(thread_id, run_id):
#     url = f"https://api.minimax.chat/v1/threads/run/retrieve?GroupId={GroupId}"
#     payload = json.dumps({
#         "thread_id": str(thread_id),
#         "run_id": str(run_id)
#     })
#     completed = False
#     while not completed:
#         response = requests.request("GET", url, headers=headers, data=payload)
#         if response.status_code == 200:
#             response_data = response.json()
#             status = response_data.get('status', '')
#             print(f"Status: {status}")
#             if status == 'completed':
#                 completed = True
#                 print("Process completed, exiting loop.")
#             else:
#                 time.sleep(2)  # 如果状态不是completed，等待两秒后重新请求
#         else:
#             print(f"Error: {response.status_code}")
#             break  # 如果请求失败，退出循环
#     return completed
# # 主流程
# def main():
#     # 上传文档
#     file_response = create_file()
#     file_id = file_response.get('file', {}).get('file_id')
#     print("file_id:",file_id)

#     # 创建助手
#     assistant_response = create_assistant(file_id)
#     assistant_id = assistant_response.get('id', '')
#     print("assistant_id:",assistant_id)

#     # 创建线程
#     thread_response = create_thread()
#     thread_id = thread_response.get('id', '')
#     print("thread_id:",thread_id)

#     # 往线程里添加信息
#     add_message_to_thread(thread_id)  # 不保存返回值

#     # 使用助手处理该线程
#     run_response = run_thread_with_assistant(thread_id, assistant_id)
#     run_id = run_response.get('id', '')  # 假设run_response是正确的JSON响应，并包含run_id
#     print("run_id:",run_id)

#     # 检查助手处理状态
#     if check_thread_run_status(thread_id, run_id):
#         # 获取线程中助手处理出的新信息
#         thread_messages_response = get_thread_messages(thread_id)
#         # 打印JSON数据
#         print(json.dumps(thread_messages_response, indent=4, ensure_ascii=False))

# if __name__ == "__main__":
#     main()
